import numpy as np

import tvm
from tvm import te
from tvm import rpc
from tvm.contrib import utils

print("run opencl test on rk3399")

n = tvm.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
s = te.create_schedule(B.op)

# NOTE: This is the setting for my rk3399 board. You need to modify
# them according to your environment.
target_host = "llvm -mtriple=aarch64-linux-gnu"
opencl_device_host = "192.168.100.249"
opencl_device_port = 9090

# create schedule for the above "add one" compute declaration
s = te.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], factor=32)
s[B].bind(xo, te.thread_axis("blockIdx.x"))
s[B].bind(xi, te.thread_axis("threadIdx.x"))
func = tvm.build(s, [A, B], "opencl", target_host=target_host)

remote = rpc.connect(opencl_device_host, opencl_device_port)

# export and upload
temp = utils.tempdir()
path = temp.relpath("lib_cl.tar")
func.export_library(path)
remote.upload(path)
func = remote.load_module("lib_cl.tar")

# run
ctx = remote.cl()
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
func(a, b)
np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)
print("OpenCL test passed!")
